import torch
from torch import nn
from d2l import torch as d2l

def corr2(X,K):
    h,w=K.shape
    Y=torch.zeros((X.shape[0]-h+1,X.shape[1]-w+1))    
    for i in range(Y.shape[0]):
        for j in range(Y.shape[1]):
            Y[i,j]=(X[i:i+h,j:j+w]*K).sum()
    return Y


class Conv2D(nn.Module):
    def __init__(self,kernel_size):
        super().__init__()
        self.weight = nn.parameter(torch.rand(kernel_size))
        self.bias = nn.Parameter(torch.zeros(1))
    def forward(self,x):
        return corr2(x,self.weight)+self.bias
# 解释一下parameter的作用，作用是在于将参数绑定在整个模型之中，以便每次更新优化参数
# 其将其他类型转化为了参数类型，方便模型去进行理解

# 实验一：图像中目标的边缘检测
# 原理：将卷积核设置为1和-1并且宽度设置为2，这样就可进行检测，因为求和的原理

X = torch.ones((6,8))
X[:,2:6] = 0
K = torch.tensor([[1.0, -1.0]])
Y=corr2(X,K)

# 实验二：进行卷积学习
# 内置的二维卷积层
conv2d=nn.Conv2d(1,1,kernel_size=(1,2),bias=False)
# 这个二维卷积层使用四维输入和输出格式（批量大小、通道、高度、宽度），
# 批量大小和通道都为一
X= X.reshape((1,1,6,8))
Y=Y.reshape((1,1,6,7))
lr=3e-2
for i in range(10):
    Y_hat = conv2d(X)
    l=(Y_hat-Y)**2
    conv2d.zero_grad()
    l.sum().backward()
    conv2d.weight.data[:]-=lr*conv2d.weight.grad
    print("epoch%d,loss:%d",i+1,l.sum())
print(conv2d.weight.data.reshape(1,2))